#include "nn.h"

void data_matrix::matmul(const data_linear& vect, data_linear& dst) const {
    if(vect.length != width) throw std::runtime_error("vect length does not match width");
    if(dst.length != height) throw std::runtime_error("dst length does not match width");
    if(vector == nullptr) throw std::runtime_error("matmul with not allocated matrix");
    if(vect.vector == nullptr) throw std::runtime_error("matmul with not allocated vector");
    if(dst.vector == nullptr) dst.allocate();
    dst.set0();
    for(int i = 0; i < height; i ++) {
        for(int j = 0; j < width; j ++)
            dst[i] = get(i, j) * vect[j];
    }
}

void grad::combine(const grad& outs_grad, grad& dst) const {
    if(outs_grad.in_size != out_size) throw std::runtime_error("Grads do not match");
    if(dst.in_size != in_size || dst.out_size != outs_grad.out_size) throw std::runtime_error("Grads do not match");
    if(!is_allocated()) throw std::runtime_error("combine of not allocated grad");
    if(!outs_grad.is_allocated()) throw std::runtime_error("combine with not allocated grad");
    if(!dst.is_allocated()) dst.allocate();
    _combine(outs_grad, dst);
}

grad* scalar_grad::get_combine_dst(const grad& outs_grad) const {
    if(outs_grad.in_size != out_size) throw std::runtime_error("scalar_grad get_combine_dst: Grads do not match");
    if(outs_grad.type == grad_type::DENSE) return new dense_grad(in_size, outs_grad.out_size);
    if(outs_grad.type == grad_type::SCALAR) return new scalar_grad(in_size);
    if(outs_grad.type == grad_type::LINCOMB) return new lincomb_modified_grad(((lincomb_grad&)outs_grad).vars, outs_grad.out_size);
    if(outs_grad.type == grad_type::LINCOMB_MODIFIED) return new lincomb_modified_grad(((lincomb_modified_grad&)outs_grad).vars, outs_grad.out_size);
    return new dense_grad(in_size, outs_grad.out_size);
}

void scalar_grad::_combine(const grad& _outs_grad, grad& _dst) const {
    if(_outs_grad.type == grad_type::DENSE) {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("scalar_grad * dense_grad: Wrong dst type");
        dense_grad& outs_grad = (dense_grad&)_outs_grad;
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < dst.out_size; i ++)
            for(int j = 0; j < in_size; j ++)
                dst.vector[i * in_size + j] = vector[j] * outs_grad.vector[i * out_size + j];
    } else if(_outs_grad.type == grad_type::SCALAR) {
        if(_dst.type != grad_type::SCALAR) throw std::runtime_error("scalar_grad * scalar_grad: Wrong dst type");
        scalar_grad& outs_grad = (scalar_grad&)_outs_grad;
        scalar_grad& dst = (scalar_grad&)_dst;
        for(int i = 0; i < in_size; i ++)
            dst.vector[i] = vector[i] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::LINCOMB) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("scalar_grad * lincomb_grad: Wrong dst type");
        lincomb_grad& outs_grad = (lincomb_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int j = 0; j < outs_grad.vars; j ++)
                dst.vector[i * outs_grad.vars + j] = vector[i * outs_grad.vars + j] * outs_grad.vector[j];
    } else if(_outs_grad.type == grad_type::LINCOMB_MODIFIED) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("scalar_grad * lincomb_modified_grad: Wrong dst type");
        lincomb_modified_grad& outs_grad = (lincomb_modified_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < in_size; i ++)
            dst.vector[i] = vector[i] * outs_grad.vector[i];
    } else {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("scalar_grad * any grad: Wrong dst type");
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < dst.out_size; i ++)
            for(int j = 0; j < in_size; j ++)
                dst.vector[i * in_size + j] = vector[j] * _outs_grad.get(j, i);
    }
}

dense_grad* dense_grad::get_combine_dst(const grad& outs_grad) const {
    if(outs_grad.in_size != out_size) throw std::runtime_error("dense_grad get_combine_dst: Grads do not match");
    return new dense_grad(in_size, outs_grad.out_size);
}

void dense_grad::_combine(const grad& _outs_grad, grad& _dst) const {
    if(_dst.type != grad_type::DENSE) throw std::runtime_error("dense_grad * any grad: Wrong dst type");
    dense_grad& dst = (dense_grad&)_dst;
    if(_outs_grad.type == grad_type::DENSE) {
        dense_grad& outs_grad = (dense_grad&)_outs_grad;
        memset(dst.vector, '\0', sizeof(dtype) * dst.in_size * dst.out_size);
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int k = 0; k < outs_grad.in_size; k ++)
                for(int j = 0; j < in_size; j ++)
                    dst.vector[i * in_size + j] += vector[k * in_size + j] * outs_grad.vector[i * out_size + k];
    } else if(_outs_grad.type == grad_type::SCALAR) {
        scalar_grad& outs_grad = (scalar_grad&)_outs_grad;
        for(int i = 0; i < out_size; i ++)
            for(int j = 0; j < in_size; j ++)
                dst.vector[i * in_size + j] = vector[i * in_size + j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::LINCOMB) {
        lincomb_grad& outs_grad = (lincomb_grad&)_outs_grad;
        memset(dst.vector, '\0', sizeof(dtype) * dst.in_size * dst.out_size);
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int j = 0; j < outs_grad.vars; j ++)
                for(int k = 0; k < in_size; k ++)
                    dst.vector[i * in_size + k] += vector[i * in_size * outs_grad.vars + j * in_size + k] * outs_grad.vector[j];
    } else if(_outs_grad.type == grad_type::LINCOMB_MODIFIED) {
        lincomb_modified_grad& outs_grad = (lincomb_modified_grad&)_outs_grad;
        memset(dst.vector, '\0', sizeof(dtype) * dst.in_size * dst.out_size);
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int j = 0; j < outs_grad.vars; j ++)
                for(int k = 0; k < in_size; k ++)
                    dst.vector[i * in_size + k] += vector[i * in_size * outs_grad.vars + j * in_size + k] * outs_grad.vector[i * outs_grad.vars + j];
    } else {
        memset(dst.vector, '\0', sizeof(dtype) * dst.in_size * dst.out_size);
        for(int i = 0; i < _outs_grad.out_size; i ++)
            for(int j = 0; j < _outs_grad.in_size; j ++)
                for(int k = 0; k < in_size; k ++)
                    dst.vector[i * in_size + k] += vector[k * in_size + k] * _outs_grad.get(j, i);
    }
}

grad* lincomb_grad::get_combine_dst(const grad& outs_grad) const {
    if(outs_grad.in_size != out_size) throw std::runtime_error("lincomb_grad get_combine_dst: Grads do not match");
    if(outs_grad.type == grad_type::SCALAR) return new lincomb_modified_grad(vars, out_size);
    if(outs_grad.type == grad_type::LINCOMB) return new lincomb_grad(vars * ((lincomb_grad&)outs_grad).vars, outs_grad.out_size);
    if(outs_grad.type == grad_type::LINCOMB_MODIFIED) return new lincomb_modified_grad(vars * ((lincomb_modified_grad&)outs_grad).vars, outs_grad.out_size);
    return new dense_grad(in_size, outs_grad.out_size);
}

void lincomb_grad::_combine(const grad& _outs_grad, grad& _dst) const {
    if(_outs_grad.type == grad_type::SCALAR) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("lincomb_grad * scalar_grad: Wrong dst type");
        scalar_grad& outs_grad = (scalar_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < out_size; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::LINCOMB) {
        if(_dst.type != grad_type::LINCOMB) throw std::runtime_error("lincomb_grad * lincomb_grad: Wrong dst type");
        lincomb_grad& outs_grad = (lincomb_grad&)_outs_grad;
        lincomb_grad& dst = (lincomb_grad&)_dst;
        for(int i = 0; i < outs_grad.vars; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::LINCOMB_MODIFIED) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("lincomb_grad * lincomb_modified_grad: Wrong dst type");
        lincomb_modified_grad& outs_grad = (lincomb_modified_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < out_size; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::DENSE) {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("lincomb_grad * dense_grad: Wrong dst type");
        dense_grad& outs_grad = (dense_grad&)_outs_grad;
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < outs_grad.out_size * outs_grad.in_size; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[j] * outs_grad.vector[i];
    } else {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("lincomb_grad * any grad: Wrong dst type");
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < _outs_grad.out_size; i ++)
            for(int j = 0; j < _outs_grad.in_size; j ++)
                for(int k = 0; k < vars; k ++)
                    dst.vector[i * in_size + j * vars + k] = vector[k] * _outs_grad.get(j, i);
    }
}

grad* lincomb_modified_grad::get_combine_dst(const grad& outs_grad) const {
    if(outs_grad.in_size != out_size) throw std::runtime_error("lincomb_modified_grad get_combine_dst: Grads do not match");
    if(outs_grad.type == grad_type::SCALAR) return new lincomb_modified_grad(vars, out_size);
    if(outs_grad.type == grad_type::LINCOMB) return new lincomb_modified_grad(vars * ((lincomb_grad&)outs_grad).vars, outs_grad.out_size);
    if(outs_grad.type == grad_type::LINCOMB_MODIFIED) return new lincomb_modified_grad(vars * ((lincomb_modified_grad&)outs_grad).vars, outs_grad.out_size);
    return new dense_grad(in_size, outs_grad.out_size);
}

void lincomb_modified_grad::_combine(const grad& _outs_grad, grad& _dst) const {
    if(_outs_grad.type == grad_type::SCALAR) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("lincomb_modified_grad * scalar_grad: Wrong dst type");
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        scalar_grad& outs_grad = (scalar_grad&)_outs_grad;
        for(int i = 0; i < out_size; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[i * vars + j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::LINCOMB) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("lincomb_modified_grad * lincomb_grad: Wrong dst type");
        lincomb_grad& outs_grad = (lincomb_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int j = 0; j < outs_grad.in_size; j ++)
                for(int k = 0; k < vars; k ++)
                    dst.vector[i * dst.vars + j * vars + k] = vector[i * dst.vars + j * vars + k] * outs_grad.vector[j];
    } else if(_outs_grad.type == grad_type::LINCOMB_MODIFIED) {
        if(_dst.type != grad_type::LINCOMB_MODIFIED) throw std::runtime_error("lincomb_modified_grad * lincomb_modified_grad: Wrong dst type");
        lincomb_modified_grad& outs_grad = (lincomb_modified_grad&)_outs_grad;
        lincomb_modified_grad& dst = (lincomb_modified_grad&)_dst;
        for(int i = 0; i < out_size; i ++)
            for(int j = 0; j < vars; j ++)
                dst.vector[i * vars + j] = vector[i * vars + j] * outs_grad.vector[i];
    } else if(_outs_grad.type == grad_type::DENSE) {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("lincomb_modified_grad * dense_grad: Wrong dst type");
        dense_grad& outs_grad = (dense_grad&)_outs_grad;
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < outs_grad.out_size; i ++)
            for(int j = 0; j < outs_grad.in_size; j ++)
                for(int k = 0; k < vars; k ++)
                    dst.vector[i * in_size + j * vars + k] = vector[j * vars + k] * outs_grad.vector[i * outs_grad.in_size + j];
    } else {
        if(_dst.type != grad_type::DENSE) throw std::runtime_error("Wrong dst type");
        dense_grad& dst = (dense_grad&)_dst;
        for(int i = 0; i < _outs_grad.out_size; i ++)
            for(int j = 0; j < _outs_grad.in_size; j ++)
                for(int k = 0; k < vars; k ++)
                    dst.vector[i * in_size + j * vars + k] = vector[j * vars + k] * _outs_grad.get(j, i);
    }
}
